#include <unistd.h>
#include <stdlib.h>
#include <stdint.h>
#include <time.h>
#include <math.h>

#include <iostream>
#include <functional>
#include <stdexcept>
#include <string>
#include <string_view>
#include <mutex>
#include <shared_mutex>

#include <libpmemobj.h>

#include <libpmemobj++/allocator.hpp>
#include <libpmemobj++/experimental/inline_string.hpp>
#include <libpmemobj++/container/array.hpp>
#include <libpmemobj++/container/string.hpp>
#include <libpmemobj++/container/vector.hpp>
#include <libpmemobj++/container/basic_string.hpp>
#include <libpmemobj++/p.hpp>
#include <libpmemobj++/persistent_ptr.hpp>
#include <libpmemobj++/pext.hpp>
#include <libpmemobj++/pool.hpp>
#include <libpmemobj++/transaction.hpp>
#include <libpmemobj++/utils.hpp>

namespace PMKV {

// **************************************** 
// define TEST to enable LOG_INFO
#define TEST
#undef TEST

#ifdef TEST
#define ENABLE_LOG true
#else
#define ENABLE_LOG false
#endif

#define LOG_WITH_FUNC __FILE__ << " " << __LINE__ << " " << __FUNCTION__ << "\n  "
#define LOG_INFO() !ENABLE_LOG ? (void)0 : Vodify() & std::cout << "LOG_INFO: "
#define LOG_ERR()  !ENABLE_LOG ? (void)0 : Vodify() & std::cerr << "LOG_ERR: "

class Vodify {
public:
    void operator&(const std::ostream&) {}
};

// **************************************** 
// PMEMDevice

const int N = 1 * 1024 * 1024;
const int LOCK_SEGMENTS = 10;
const int KEY_MAXSIZE = 256;

struct element {
    char key[256] = {0};
    pmem::obj::p<uint64_t> _size;   ///< val valid data size
    /**
     * @todo PERF splitting metadata and value data causes multiple roundtrips
     * when updating a kv entry, since both fields need to be written.
     */
    pmem::obj::p<double> _mtime;    ///< modified time
    pmem::obj::string val;          ///< user data

    uint64_t size() const {
        LOG_INFO() << "element::size()= " << _size << std::endl;
        return _size;
    }

    double mtime() const {
        LOG_INFO() << "element::touch()=" << _mtime << std::endl;
        return _mtime;
    }

    char* data() {
        LOG_INFO() << "element::data()" << std::endl;
        return val.data();
    }

    const char *data() const {
        LOG_INFO() << "element::data() const" << std::endl;
        return val.data();
    }

    int read(uint64_t off, uint64_t len, char* buf) const {
        LOG_INFO() << "element::read(off, len, buf) const: " << off << " " << len << " " << buf << std::endl;
        // 读全部有效数据
        if(len == 0) {off = 0;len = _size;}
        else if(off + len > _size) len = _size - off;

        // read val, [off, off + len]
        const char* data = val.data();
        memcpy(buf, data + off, len);
        return len;
    }

    // attention: call it in pmem::obj::transaction::run(pop, [&]{});
    int resize(uint64_t size) {
        LOG_INFO() << "element::resize(): " << size << std::endl;
        if(size > val.capacity()) val.reserve(size);
        _size = size;
        return 0;
    }

    // attention: call it in pmem::obj::transaction::run(pop, [&]{});
    int touch(double mtime) {
        LOG_INFO() << "element::touch(mtime): " << mtime << std::endl;
        _mtime = mtime;
        return 0;
    }

    // attention: call it in pmem::obj::transaction::run(pop, [&]{});
    char* range(uint64_t off, uint64_t len) {
        LOG_INFO() << "element::range(off, len): " << off << " " << len << std::endl;
        // 越界，重分配
        if(off + len > val.capacity()) val.reserve(off + len);
        // 超过有效数据
        if(off + len > _size) _size = off + len;

        const char *addr = val.cdata() + off;
        pmemobj_tx_add_range_direct(addr, len);
        return const_cast<char*>(addr);
    }

    // attention: call it in pmem::obj::transaction::run(pop, [&]{});
    int zero(uint64_t off, uint64_t len) {
        LOG_INFO() << "element::zero(off, len): " << off << " " << len << std::endl;
        // 越界，重分配
        if(off + len > val.capacity()) val.reserve(off + len);
        // 超过有效数据
        if(off + len > _size) _size = off + len; 

        // zero val, [off, off + len]
        char* data = val.data();
        TX_MEMSET(data + off, 0, len);
        return len;
    }

    // attention: call it in pmem::obj::transaction::run(pop, [&]{});
    int write(uint64_t off, uint64_t len, char* buf) {
        LOG_INFO() << "element::write(off, len, buf): " << off << " " << len << " " << buf << std::endl;
        // 越界，重分配
        // if fix the max size, return EFBIG
        if(off + len > val.capacity()) val.reserve(off + len);
        // 超过有效数据
        if(off + len > _size) _size = off + len;

        // write val, [off, off + len]
        const char *data = val.cdata();
        TX_MEMCPY(const_cast<char*>(data) + off, buf, len);
        return len;
    }
};

/**
 * iterator to elements
 * @note invalidated on element creation and removal
 * @tparam E 
 */
template<typename E>
struct iterator_impl {
    using self_type = iterator_impl<E>;

    uint64_t index;
    uint64_t eindex;
    E* ptr = nullptr;

    E* operator->() const {
        return ptr;
    }

    operator bool() const{
        return ptr != nullptr;
    }

    bool operator==(const self_type& rhs) const {
        return ptr == rhs.ptr;
    }

    bool operator!=(const self_type& rhs) const {
        return ptr != rhs.ptr;
    }

    bool empty() {
        if(ptr == nullptr) return true;
        return false;
    }
};
using iterator = iterator_impl<element>;
using const_iterator = iterator_impl<const element>;

using element_type = pmem::obj::persistent_ptr<element>;
using bucket_type = pmem::obj::vector<element_type>;
using bucket_array_type = pmem::obj::array<bucket_type, N>;

struct kv_type {
    bucket_array_type buckets;
};

struct root {
    pmem::obj::persistent_ptr<kv_type> kv;
};

class PMEMDevice {
private:
    pmem::obj::persistent_ptr<root> r;
    std::array<std::shared_mutex, LOCK_SEGMENTS> seg_locks;
public:
    pmem::obj::pool<root> pop;

    PMEMDevice() {

    }

    explicit PMEMDevice(const char* path) {
        open(path);
    }

    ~PMEMDevice() {
        close();
    }
    
    // open pmempool
    // pmempool create obj --layout=PMEMDevice -s 40G /data/pmem0/PMEMDevice
    int open(const char *path) {
        LOG_INFO() << "open(path): " << path << std::endl;
        try {
            // pop = pmem::obj::pool<root>::create(path, "PMEMDevice", pool_size, S_IWUSR | S_IRUSR);
            pop = pmem::obj::pool<root>::open(path, "PMEMDevice");
            r = pop.root();
            if(r->kv == nullptr) {
                pmem::obj::transaction::run(pop, [&] {
                    r->kv = pmem::obj::make_persistent<kv_type>();
                });
            }
        } catch (pmem::pool_error &e) {
            std::cerr << e.what() << std::endl;
            throw std::runtime_error("PMEMDevice::open failed");
            return -1;
        } catch (std::exception &e) {
            std::cerr << e.what() << std::endl;
            return -1;
        }
        return 0;
    }

    // close pmmepool
    void close() {
        LOG_INFO() << "close()" << std::endl;
        try {
            pop.close();
        } catch (const std::logic_error &e) {
		    std::cerr << e.what() << std::endl;
	    }
    }

    pmem::obj::pool_base& get_pool() {
        // have to check valid?
        return pop;
    }

    /**
     * @brief: create a element and reserve the val to size
     * @param key
     * @param size size hint of the KV pair to be created
     * @return iterator
     */
    iterator create(const std::string& key, uint64_t size) {
        LOG_INFO() << "create(key, size): " << key << " " << size << std::endl;
        iterator it;

        auto index = std::hash<std::string>{}(key) % N;
        it.index = index;

        /* sanity check */
        const char* _key = key.c_str();
        uint32_t _len = key.size();
        if(_len > KEY_MAXSIZE) {
            throw std::length_error("create: key size too long");
        }

        /* check if already exist */
        int m;  // hash bucket index
        {
            const auto &bucket_cref = r->kv->buckets.const_at(index);
            m = bucket_cref.size();
            for(int i = 0; i < m; ++i) {
                if(strcmp(bucket_cref[i]->key, _key) == 0) {
                    it.eindex = i;
                    it.ptr = bucket_cref[i].get();

                    LOG_INFO()
                    << "create " << key
                    << ": already exist!" << std::endl;
                    return it;
                }
            }
        }

        /* actually insert */
        try {
        pmem::obj::flat_transaction::run(pop, [&]{
            auto &bucket_wref = r->kv->buckets.at(index);
            auto &inserted_wref = bucket_wref.emplace_back(pmem::obj::make_persistent<element>());
            it.eindex = m;
            it.ptr = inserted_wref.get();
            TX_MEMCPY(it->key, _key, _len);
            it->val.reserve(size);
            it->_size = 0;
            LOG_INFO() << "it->val.capacity: " << it->val.capacity() << std::endl;
        });
        }
        catch(pmem::transaction_out_of_memory &e) {
            throw std::out_of_range("create: pmem no space allocate");
        }
        LOG_INFO()
        << "create " << key 
        << ": succeed!" << std::endl;
        return it;
    }

    /**
     * get mutex of the corresponding lock segment
     * @param key 
     * @return 
     */
    std::shared_mutex &get_hoid_segment_mutex(std::string_view key) {
        auto index = std::hash<std::string_view>{}(key) % LOCK_SEGMENTS;
        return seg_locks[index];
    }

    /**
     * @brief: find an element
     * @param key
     * @return iterator
     */
    iterator find(const std::string &key) const {
        LOG_INFO() << "find(key): " << key << std::endl;
        iterator it;

        auto index = std::hash<std::string>{}(key) % N;
        it.index = index;

        const char* _key = key.data();

        const auto &bucket_cref = r->kv->buckets.const_at(index);
        const int m = bucket_cref.size();
        for(int i = 0; i < m; ++i) {
            if(strcmp(bucket_cref[i]->key, _key) == 0) {
                it.eindex = i;
                it.ptr = bucket_cref[i].get();
                LOG_INFO()
                << "find " << key
                << ": index " << it.index << "-" << it.eindex << std::endl;
                return it;
            }
        }

        LOG_INFO()
        << "find " << key << " in bucket " << it.index
        << ": not found!" << std::endl;
        return it;
    }

    /**
     * @brief: find a const element 
     * @note this is essentially the same as the non-const find(), differs in
     *      that this method returns a const iterator, only for read-only use,
     *      write (and for that matter, snapshots) are disabled statically so
     *      you can't make mistakes.
     * @param key
     * @return const iterator
     */
    const_iterator cfind(const std::string key) const {
        LOG_INFO() << "cfind(key): " << key << std::endl;
        const_iterator it;

        auto index = std::hash<std::string>{}(key) % N;
        it.index = index;
        
        const char* _key = key.c_str();

        const auto &bucket_cref = r->kv->buckets.const_at(index);
        const int m = bucket_cref.size();
        for(int i = 0; i < m; ++i) {
            if(strcmp(bucket_cref[i]->key, _key) == 0) {
                it.eindex = i;
                it.ptr = bucket_cref[i].get();
                LOG_INFO()
                << "find " << key
                << ": index " << it.index << "-" << it.eindex << std::endl;
                return it;
            }
        }
        
        LOG_INFO()
        << "find " << key << " in bucket " << it.index
        << ": not found!" << std::endl;
        return it;
    }

    iterator end() const {
        LOG_INFO() << "end()" << std::endl;
        iterator it;
        return it;
    }

    const_iterator cend() const {
        LOG_INFO() << "cend()" << std::endl;
        const_iterator it;
        return it;
    }

    /**
     * @brief: free space, remove the element
     * @todo PERF (zxg) eagerly erase remove element from dataset, causing
     *      frequentddata shifting (and snapshot along)
     * @param key
     * @return -1/fail
     */
    int remove(const std::string &key) {
        LOG_INFO() << "remove(key): " << key << std::endl;
        auto index = std::hash<std::string>{}(key) % N;
        const char* _key = key.c_str();

        const auto &bucket_cref = r->kv->buckets.const_at(index);
        const int m = bucket_cref.size();
        for(int i = 0; i < m; ++i){
            if(strcmp(bucket_cref[i]->key, _key) == 0){
                pmem::obj::flat_transaction::run(pop, [&]{
                    auto &bucket_wref = r->kv->buckets.at(index);   // writable ref to hash bucket
                    // auto &pptrkv_wref = bucket_wref.at(i);          // writable ref to persistent ptr to kv pair
                    // pmem::obj::delete_persistent<element>(pptrkv_wref);
                    // pptrkv_wref = nullptr;
                    pmem::obj::delete_persistent<element>(bucket_wref.const_at(i));
                    bucket_wref.erase(bucket_wref.cbegin() + i);
                });

                LOG_INFO()
                << "remove " << key
                << ": succeed!" << std::endl;  
                return 0;
            }
        }
        
        LOG_ERR() 
        << "remove " << key
        << ": out of range!" << std::endl;
        // throw std::out_of_range("remove: no entry in hash_map");
        return -1;
    }

    /**
     * @brief: free space, remove the element, invalidating ALL existing iterators
     * @param iterator
     * @return 0
     */
    template<typename E>
    int remove(iterator_impl<E>& it) {
        if(it.ptr == nullptr) return -1;
        LOG_INFO() << "remove(iterator): key " << it->key 
        << " index " << it.index << "-" << it.eindex << std::endl;
        it.ptr = nullptr;
        pmem::obj::flat_transaction::run(pop, [&]{
            auto &bucket_wref = r->kv->buckets.at(it.index);
            pmem::obj::delete_persistent<element>(bucket_wref.const_at(it.index));
            // bucket_wref[it.eindex] = nullptr;
            bucket_wref.erase(bucket_wref.cbegin() + it.eindex);
        });
        LOG_INFO()
        << "remove : succeed!" << std::endl;
        return 0;
    }

    /**
     * @brief: write buf[off, off + len] to val
     * @param {string&} key
     * @param {uint64_t} off
     * @param {uint64_t} len
     * @param {char} *buf
     * @return succeed write bytes
     */    
    int write(const std::string& key, uint64_t off, uint64_t len, const char *buf) {
        LOG_INFO() << "write(key, off, len, buf): " 
        << key << " " << off << " " << len << " " << buf << std::endl;
        auto index = std::hash<std::string>{}(key) % N;
        const char* _key = key.c_str();

        for(const auto &e : r->kv->buckets.const_at(index)) {
            if (strcmp(e->key, _key) == 0) {
                pmem::obj::transaction::run(pop, [&] {
                    // pmem::obj::transaction::snapshot(&v)
                    const char* data = e->val.cdata();
                    TX_MEMCPY(const_cast<char*>(data) + off, buf, len);
                });

                LOG_INFO()
                << "write " << e->key << "-" << e->val.data()
                << ": succeed!" << std::endl;
                return len;
            }
        }
        LOG_ERR()
        << "write"  << key
        << ": out of range!" << std::endl;
        throw std::out_of_range("no entry in hash_map");
        return -1;
    }

    // read len bytes from key-value (data + off) to buf
    int read(const std::string& key, uint64_t off, uint64_t len, char *buf) const {
        LOG_INFO() << "read(key, off, len, buf): " 
        << key << " " << off << " " << len << " " << buf << std::endl;
        auto index = std::hash<std::string>{}(key) % N;
        const char* _key = key.c_str();

        for (const auto &e : r->kv->buckets.const_at(index)) {
			if (strcmp(e->key, _key) == 0) {
                const char* data = e->val.cdata();
                memcpy(buf, data + off, len);
                LOG_INFO() 
                << "read " << e->key << "-" << e->val.data()
                << ": succeed!" << std::endl;
                return len;
            }
		}

        LOG_ERR()
        << "read " << key
        << ": out of range!" << std::endl;
        throw std::out_of_range("no entry in hash_map");
        return -1;
    }

    int touch(const std::string& key, double mtime) {
        auto index = std::hash<std::string>{}(key) % N;
        const char* _key = key.c_str();
        for(const auto &e : r->kv->buckets.const_at(index)) {
            if (strcmp(e->key, _key) == 0) {
                pmem::obj::flat_transaction::run(pop, [&] {
                    e->_mtime = mtime;
                });
                LOG_INFO()
                << "touch " << key << " at " << mtime
                << ": succeed!" << std::endl;
                return 0;
            }
        }

        LOG_ERR()
        << "touch " << key
        << ": out of range!" << std::endl;
        throw std::out_of_range("no entry in hash_map");
        return -1;
    }
};

} // namespace PMKV